{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lTx8eQlc3cP-"
      },
      "source": [
        "##### Copyright 2024 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "4HZoi8yf4GEU"
      },
      "outputs": [],
      "source": [
        "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M9I7LG483nXB"
      },
      "source": [
        "# Gemini API: LLM Observability with MLflow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "awKO767lQIWh"
      },
      "source": [
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/cookbook/blob/main/examples/mlflow/MLflow_Observability.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" height=30/></a>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b1xoF_bU4NCP"
      },
      "source": [
        "## Overview"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CWedABji6bXJ"
      },
      "source": [
        "[MLflow](https://mlflow.org/) is an open-source platform to assist machine learning practitioners and teams in handling the complexities of the machine learning process.\n",
        "\n",
        "It provides [MLflow Tracing](https://mlflow.org/docs/latest/tracing/) that enhances LLM observability in your Generative AI applications by capturing detailed information about the execution of your application’s services. Tracing provides a way to record the inputs, outputs, and metadata associated with each intermediate step of a request, enabling you to easily pinpoint the source of bugs and unexpected behaviors.\n",
        "\n",
        "MLflow provides a built-in integration with Google Gen AI SDK that enables you to instrument your Gemini calls easily. This cookbook describes the basic usage of the MLflow tracing integration with the `google-genai` package."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "baa4fac6585c"
      },
      "source": [
        "<!-- Community Contributor Badge -->\n",
        "<table>\n",
        "  <tr>\n",
        "    <!-- Author Avatar Cell -->\n",
        "    <td bgcolor=\"#d7e6ff\">\n",
        "      <a href=\"https://github.com/TomeHirata\" target=\"_blank\" title=\"View Tomu Hirata's profile on GitHub\">\n",
        "        <img src=\"https://github.com/TomeHirata.png?size=100\"\n",
        "             alt=\"TomeHirata's GitHub avatar\"\n",
        "             width=\"100\"\n",
        "             height=\"100\">\n",
        "      </a>\n",
        "    </td>\n",
        "    <!-- Text Content Cell -->\n",
        "    <td bgcolor=\"#d7e6ff\">\n",
        "      <h2><font color='black'>This notebook was contributed by <a href=\"https://github.com/TomeHirata\" target=\"_blank\"><font color='#217bfe'><strong>Tomu Hirata</strong></font></a>.</font></h2>\n",
        "      <!-- Footer -->\n",
        "      <font color='black'><small><em>Have a cool Gemini example? Feel free to <a href=\"https://github.com/google-gemini/cookbook/blob/main/CONTRIBUTING.md\" target=\"_blank\"><font color=\"#078efb\">share it too</font></a>!</em></small></font>\n",
        "    </td>\n",
        "  </tr>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "70wYOKUC8q1m"
      },
      "source": [
        "## Installation\n",
        "Install Google Gen AI SDK (`google-genai`) and MLflow (`mlflow`). See [troubleshooting for MLFlow installation](https://mlflow.org/docs/latest/quickstart_drilldown/#quickstart_drilldown_install) for other install options."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "mnQbBnA1GKha"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.0.1\u001b[0m\n",
            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
            "Note: you may need to restart the kernel to use updated packages.\n"
          ]
        }
      ],
      "source": [
        "# Integration with google-genai is supported mlflow >= 2.20.3\n",
        "%pip install -q google-genai \"mlflow>=2.20.3\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eomJzCa6lb90"
      },
      "source": [
        "## Create Gemini Client with your API key\n",
        "\n",
        "Let's create an API client and pass your API key. If you do not have API ket yet, visit [AI Studio](https://aistudio.google.com/app/apikey) to create one."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "v-JZzORUpVR2"
      },
      "outputs": [],
      "source": [
        "import google.genai as genai\n",
        "from google.colab import userdata\n",
        "\n",
        "client = genai.Client(api_key=userdata.get('GOOGLE_API_KEY'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "36730d1ed562"
      },
      "source": [
        "## Tracking Server\n",
        "There are several options to run MLflow tracking server: local tracking server, Databricks Free Trial, and production Databricks. See [our documentation](https://mlflow.org/docs/latest/getting-started/tracking-server-overview/index.html) for the comparison. In this example, you use [Databricks Free Trial](https://mlflow.org/docs/latest/getting-started/databricks-trial/) which allows easy connection from Colab notebooks and enable you to use managed MLflow for free. Follow the steps below to create an account and generate a Personal Access Token (PAT) to connect to your workspace.\n",
        "\n",
        "* Go to the [Databricks Trial Signup Page](http://signup.databricks.com/) and create your account\n",
        "* Follow the steps in [this guide](https://docs.databricks.com/aws/en/dev-tools/auth/pat) to create a PAT for your Databricks workspace user\n",
        "* You need following information to connect to your workspace from your notebook\n",
        "  * Databricks Host: Use \"https://\\<your workspace host\\>.cloud.databricks.com/\n",
        "  * Token: Your personal access token for your Databricks Workspace. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "ba6aeb4716e9"
      },
      "outputs": [],
      "source": [
        "WORKSPACE_EMAIL = userdata.get(\"WORKSPACE_EMAIL\") # your email\n",
        "WORKSPACE_PAT = userdata.get(\"WORKSPACE_PAT\") # your databricks host https://<your workspace host>.cloud.databricks.com/\n",
        "WORKSPACE_URL = userdata.get(\"WORKSPACE_URL\") # your PAT"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6c0a3fb64bf3"
      },
      "source": [
        "Then, set the tracking server uri and experiment name."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d362a772c8b1"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import mlflow\n",
        "mlflow.set_tracking_uri('databricks')\n",
        "mlflow.set_registry_uri('databricks-uc')\n",
        "os.environ['DATABRICKS_HOST'] = WORKSPACE_URL\n",
        "os.environ['DATABRICKS_TOKEN'] = WORKSPACE_PAT\n",
        "mlflow.login()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "69496e82f5e0"
      },
      "outputs": [],
      "source": [
        "# Create Experiment\n",
        "mlflow.create_experiment(\n",
        "    f\"/Users/{WORKSPACE_EMAIL}/gemini\",\n",
        "    artifact_location=\"dbfs:/Volumes/workspace/default/gemini\",\n",
        ")\n",
        "mlflow.set_experiment(f\"/Users/{WORKSPACE_EMAIL}/gemini\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ca298fb29003"
      },
      "source": [
        "## Enable AutoLogging\n",
        "MLflow Tracing provides automatic tracing capability for Google Gemini. By enabling auto tracing for Gemini by calling the `mlflow.gemini.autolog()` function, MLflow will capture nested traces and log them to the active MLflow Experiment upon invocation of Gemini Python SDK.\n",
        "\n",
        "MLflow trace automatically captures the following information about Gemini calls:\n",
        "- Prompts and completion responses\n",
        "- Latencies\n",
        "- Model name\n",
        "- Additional metadata such as temperature, max_tokens, if specified.\n",
        "- Function calling if returned in the response\n",
        "- Any exception if raised"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "f4905e46ad24"
      },
      "outputs": [],
      "source": [
        "import mlflow\n",
        "\n",
        "# Turn on auto tracing for Gemini\n",
        "mlflow.gemini.autolog()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "79d326eb8531"
      },
      "source": [
        "## Call Simple Content Generation\n",
        "Let's run `client.models.generate_content` to try a simple text generation use case with MLflow tracing. For Jupyter Notebook users, MLflow provides a convenient way to see the generated traces on your notebook. See [this blog](https://mlflow.org/blog/mlflow-tracing-in-jupyter) for more information. For users who use other platforms, visit \"http://localhost:5000\" to see MLflow UI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3fde5077aa51"
      },
      "outputs": [],
      "source": [
        "MODEL_ID = \"gemini-2.5-flash\" # @param [\"gemini-2.5-flash-lite\",\"gemini-2.5-flash\",\"gemini-2.5-pro\",\"gemini-3-pro-preview\"] {\"allow-input\":true, isTemplate: true}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "e678e88a5277"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'cold\\n'"
            ]
          },
          "execution_count": 17,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "result = client.models.generate_content(\n",
        "    model=MODEL_ID, contents=\"The opposite of hot is\"\n",
        ")\n",
        "result.text"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fb73f557509b"
      },
      "source": [
        "<img src=\"https://i.imgur.com/xWGZXtZ.png\" />"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "be1d5d5dd4df"
      },
      "source": [
        "## Multi-Turn Chat Interactions\n",
        "MLflow tracing captures the structure of your interactions with the Gemini API. Run the following cell to try multi-turn chat and see how MLflow captures the interaction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "0208fc963eca"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "A computer follows instructions, like a recipe, to do things with numbers and pictures.\n",
            "\n",
            "A computer works by executing a sequence of instructions, called a program, written in a language it understands. These instructions manipulate binary data (0s and 1s) representing information, performing calculations, storing and retrieving data from memory and storage devices, and interacting with input and output devices like the keyboard, screen, and network to complete tasks.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "with mlflow.start_span(name=\"multi-turn\"):\n",
        "    chat = client.chats.create(model=MODEL_ID)\n",
        "    response = chat.send_message(\"In one sentence, explain how a computer works to a young child.\")\n",
        "    print(response.text)\n",
        "    response = chat.send_message(\"Okay, how about a more detailed explanation to a high schooler?\")\n",
        "    print(response.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "87aeb4d555a0"
      },
      "source": [
        "<img src=\"https://i.imgur.com/ftJ3CEp.png\">"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "03bef09f13d8"
      },
      "source": [
        "## Function Call\n",
        "If your application uses [function calling](https://ai.google.dev/gemini-api/docs/function-calling) of the Gemini API, the function definition, function call and function response are captured by MLflow automatically."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "e33a70113f74"
      },
      "outputs": [],
      "source": [
        "def add(a: float, b: float):\n",
        "    \"\"\"returns a + b.\"\"\"\n",
        "    return a + b\n",
        "\n",
        "\n",
        "def subtract(a: float, b: float):\n",
        "    \"\"\"returns a - b.\"\"\"\n",
        "    return a - b\n",
        "\n",
        "\n",
        "def multiply(a: float, b: float):\n",
        "    \"\"\"returns a * b.\"\"\"\n",
        "    return a * b\n",
        "\n",
        "\n",
        "def divide(a: float, b: float):\n",
        "    \"\"\"returns a / b.\"\"\"\n",
        "    return a / b\n",
        "\n",
        "\n",
        "operation_tools = [add, subtract, multiply, divide]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "c77b9d6d38d1"
      },
      "outputs": [],
      "source": [
        "chat = client.chats.create(\n",
        "    model = MODEL_ID,\n",
        "    config = {\n",
        "        \"tools\": operation_tools,\n",
        "    }\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "66c39be9edcf"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'2508\\n'"
            ]
          },
          "execution_count": 24,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "response = chat.send_message(\n",
        "    \"I have 57 cats, each owns 44 mittens, how many mittens is that in total?\"\n",
        ")\n",
        "response.text"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "249a739836e7"
      },
      "source": [
        "<img src=\"https://i.imgur.com/mJeryzg.png\"/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tt1wSSMIxsf2"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "That's all for this cookbook. MLflow tracing provides many features that are not included in this notebook and actively releases new features. Visit [MLflow Gemini tracing integration](https://mlflow.org/docs/latest/tracing/integrations/gemini) for more configurations and [MLflow Tracing Overview](https://mlflow.org/docs/latest/tracing/) for general offerings."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "MLflow_Observability.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
